package com.wm.wmim.connect.ws.manager;

import com.wm.wmim.common.constant.Constant;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.collections4.MapUtils;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;

import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.ReentrantLock;

/**
 * @author weiming
 * @date 2024/8/17 23:53
 * @description
 */
@Slf4j
public class WebSocketSessionManager {

    /**
     * userId : session
     */
    private final static Map<String, Map<Long, List<WebSocketSession>>> userToSessionMap = new ConcurrentHashMap<>();

    /**
     * 锁
     */
    private static final Map<WebSocketSession, ReentrantLock> lockMap = new ConcurrentHashMap<>();


    public static void saveSession(WebSocketSession webSocketSession) {
        if (webSocketSession == null) {
            return;
        }
        Long userId = WebSocketSessionManager.getSessionAttribute(webSocketSession, Constant.WsParamConstant.USER_ID, Long.class);
        String businessScene = WebSocketSessionManager.getSessionAttribute(webSocketSession, Constant.WsParamConstant.BUSINESS_SCENE, String.class);
        // 一个用户单场景多个session, 注意可能分布在不同机器
        userToSessionMap.computeIfAbsent(businessScene, s -> new ConcurrentHashMap<>())
                .computeIfAbsent(userId, s -> new CopyOnWriteArrayList<>()).add(webSocketSession);
    }

    public static List<WebSocketSession> getSession(String businessScene, Long userId) {
        Map<Long, List<WebSocketSession>> map = userToSessionMap.get(businessScene);
        if (map != null) {
            return map.get(userId);
        }
        return null;
    }


    public static void removeSession(WebSocketSession webSocketSession) {
        if (webSocketSession == null) {
            return;
        }
        Long userId = WebSocketSessionManager.getSessionAttribute(webSocketSession, Constant.WsParamConstant.USER_ID, Long.class);
        String businessScene = WebSocketSessionManager.getSessionAttribute(webSocketSession, Constant.WsParamConstant.BUSINESS_SCENE, String.class);
        Map<Long, List<WebSocketSession>> userSessionMap = userToSessionMap.get(businessScene);
        List<WebSocketSession> webSocketSessions = userSessionMap.get(userId);
        webSocketSessions.remove(webSocketSessions);
        // 避免MAP过大，且不能被回收
        if (CollectionUtils.isEmpty(webSocketSessions)) {
            userSessionMap.remove(userId);
        }
        if (MapUtils.isEmpty(userSessionMap)) {
            userToSessionMap.remove(businessScene);
        }
        removeLock(webSocketSession);
        try {
            webSocketSession.close();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    public static <T> T getSessionAttribute(WebSocketSession webSocketSession, String attributeCode, Class<T> attClazz) {
        Map<String, Object> attributes = webSocketSession.getAttributes();
        try {
            if (MapUtils.isNotEmpty(attributes)) {
                Object o = attributes.get(attributeCode);
                return attClazz.cast(o);
            }
        } catch (Exception e) {
            log.error("[SessionManager] 获取会话属性异常! session={},att={},clazz={}", webSocketSession, attributeCode, attClazz, e);
        }
        return null;
    }


    public static boolean sendMessage(WebSocketSession session, String message) {
        return withLock(10000, session, message, () -> doSendMessage(session, message));
    }


    private static boolean doSendMessage(WebSocketSession session, String message) {
        if (session != null && session.isOpen()) {
            try {
                session.sendMessage(new TextMessage(message));
                return true;
            } catch (Exception e) {
                log.error("[SessionManager] 发送消息异常! session={},message={}", session, message, e);
            }
        }
        return false;
    }


    private static boolean withLock(Integer waitMs, WebSocketSession session, String message, Runnable action) {
        if (session == null) {
            return false;
        }
        String sessionId = session.getId();
        // 公平锁
        ReentrantLock lock = lockMap.computeIfAbsent(session, k -> new ReentrantLock(true));
        try {
            if (lock.tryLock(waitMs, TimeUnit.MILLISECONDS)) {
                try {
                    action.run();
                    return true;
                } finally {
                    lock.unlock();
                }
            } else {
                // TODO 异步重试机制
                log.error("[SessionManager] 获取锁失败! waitMs={}, sessionId={}, message={}", waitMs, sessionId, message);
            }
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            log.error("[SessionManager] 消息发送异常! waitMs={}, sessionId={}, message={}", waitMs, sessionId, message, e);
        } catch (Exception e) {
            log.error("[SessionManager] 消息发送异常! waitMs={}, sessionId={}, message={}", waitMs, sessionId, message, e);
        }
        return false;
    }


    // 会话关闭时清理
    private static void removeLock(WebSocketSession session) {
        ReentrantLock lock = lockMap.remove(session);
        if (lock != null && lock.isHeldByCurrentThread()) {
            lock.unlock();
        }
    }

    private WebSocketSessionManager() {
    }

}
